import math
import argparse
import sys
import time
import numpy as np
import pandas as pd
from matplotlib import animation
import matplotlib.pyplot as plt
from src.motionplanner.Constrained_ILQR.scripts.arguments import add_arguments
from src.motionplanner.Constrained_ILQR.scripts.python_simulator.PolyRect import PolyRect
from src.motionplanner.Constrained_ILQR.scripts.ilqr.iLQR import iLQR
from bins import Map
from bins import vehicleParams
from tools.DrawModel import Draw

PI = math.pi

class CilqrModel:
    def __init__(self, args):
        self.args = args
        self.NPC_dict = {}
        self.map = Map.Map()
        self.vehicle = vehicleParams.VehicleParams()
        self.current_ego_state = None

        self.x_local_plan = []
        self.y_local_plan = []
        self.x_desired_plan = []
        self.y_desired_plan = []

        # 初始化 iLQR agent (规划器)
        self.navigation_agent = None
        # 创建 iLQR 对象
        self.NPC_states = [] # NPC车辆的状态序列
        self.X = None



    def load_NPC_trajectory(self, file_name):
        try:
            df = pd.read_csv(file_name)
        except FileNotFoundError:
            raise FileNotFoundError(f"CSV文件未找到，请检查文件路径: {file_name}")

        # expected_columns = ['x', 'y', 'v', 'theta']
        # if not all(col in df.columns for col in expected_columns):
        #     raise ValueError(f"CSV文件必须包含以下列: {expected_columns}")

        # 将数据转换为NumPy数组，初始形状为 (N, 4)
        npc_data = df.to_numpy()
        # 转置成 (4, N)，以符合后续代码使用
        self.NPC_states = npc_data.T[:, 0:self.args.horizon]

    def create_ilqr_agent(self):
        self.navigation_agent = iLQR(self.args, self.vehicle.get_BOX())
        self.navigation_agent.set_global_plan(self.plan_ilqr)
        pass

    def create_global_plan(self, desired_y, v = 0.0):
        # -------------------------
        # 创建并保存全局计划轨迹(简单地在 y=desired_y 处画一条横线)
        # -------------------------
        y = desired_y
        self.plan_ilqr = []

        for i in range(0, int((self.map.road_length+20) / 2)):
            self.plan_ilqr.append(np.array([i * 6, y, v]))
        self.plan_ilqr = np.array(self.plan_ilqr)

    def get_ego_states(self):
        ego_states = np.array([[self.current_ego_state[0], self.current_ego_state[1],                         0],
                               [self.current_ego_state[2],                         0,                         0],
                               [                        0,                         0, self.current_ego_state[3]],
                               [                        0,                         0,                         0],
                               [                        0,                         0,                         0]])
        return ego_states

    def run_step_ilqr(self):
        assert self.navigation_agent != None, "Navigation Agent not initialized"

        start = time.process_time()
        desired_path, local_plan, control, X = self.navigation_agent.run_step(self.get_ego_states(),
                                                                           [self.NPC_states])
        print(time.process_time() - start)
        print("Controller: Acc {} Steer: {}".format(control[0, 0], control[1, 0]))
        self.X = X
        return desired_path, local_plan, control[:, 0], X

    def set_current_state(self, x, y, h, v):
        self.current_ego_state = np.array([x, y, v, h])

    def saveTrajectory(self, filename):
        np.savetxt(filename, self.X, delimiter=",", fmt="%.2f")

    def load_restrained_vehicle(self, filename):
        df = pd.read_csv(filename)
        npc_data = df.to_numpy()
        self.navigation_agent.target_traj = npc_data.T




#
# if __name__ == "__main__":
#     argparser = argparse.ArgumentParser(description='CARLA CILQR')
#     add_arguments(argparser)
#     args = argparser.parse_args()
#     args.horizon = 40
#     args.x_min = 42.0
#     args.x_max = 1000000.0
#     args.y_min = 0.0
#     args.y_max = 3.4
#     args.v_min = 0.0
#     args.v_max = 23.0
#     args.w_state_bound = 100000000000000
#     args.timestep = 0.2
#     print("w_val", args.w_vel)
#     args.w_vel = 1
#     args.w_offset = 10
#     args.offset = 10
#     print("wacc",args.acc_limits)
#     args.acc_limits[1] = 4.0
#     args.w_yawrate = 100
#     test = CilqrModel(args)
#     test.load_NPC_trajectory("../trajectory/test1.csv")
#
#     test.set_current_state(42, 3.4 / 2, 0, 0)
#     test.create_global_plan(3.4 / 2, 70 / 3.6)
#     test.create_ilqr_agent()
#     test.load_restrained_vehicle("../trajectory/test1.csv")
#     desired_path, local_plan, control, X = test.run_step_ilqr()
#     print(X)
#     test.saveTrajectory("../trajectory/test2.csv")
#
#
#     draw = Draw()
#     draw.initRoad()
#     draw.drawTrajectory("../trajectory/test2.csv", "red")
#     draw.drawTrajectory("../trajectory/test1.csv", "black")
#     draw.drawCar(42, draw.road_width / 4, 0, 0 / 180 * math.pi, car_name="CarII")
#
#
#
#     draw.savefig("car2begin")

# if __name__ == "__main__":
#     argparser = argparse.ArgumentParser(description='CARLA CILQR')
#     add_arguments(argparser)
#     args = argparser.parse_args()
#     args.horizon = 40
#     args.x_min = 38.0
#     args.x_max = 1000000.0
#     args.y_min = 0.0
#     args.y_max = 3.4
#     args.v_min = 0.0
#     args.v_max = 23.0
#     args.w_state_bound = 100000000000
#     args.timestep = 0.2
#     print("w_val", args.w_vel)
#     args.w_vel = 0.5
#     args.w_offset = 0.5
#     args.offset = 10
#     print("wacc", args.acc_limits)
#     args.acc_limits[1] = 4.0
#     args.w_yawrate = 100
#     test = CilqrModel(args)
#     test.load_NPC_trajectory("../trajectory/test2.csv")
#
#     test.set_current_state(38, 3.4 / 2, 0, 0)
#     test.create_global_plan(3.4 / 2, 70 / 3.6)
#     test.create_ilqr_agent()
#     test.load_restrained_vehicle("../trajectory/test2.csv")
#     desired_path, local_plan, control, X = test.run_step_ilqr()
#     print(X)
#     test.saveTrajectory("../trajectory/test3.csv")
#
#     draw = Draw()
#     draw.initRoad()
#     draw.drawTrajectory("../trajectory/test3.csv", "red")
#     draw.drawTrajectory("../trajectory/test2.csv", "black")
#     draw.drawCar(38, draw.road_width / 4, 0, 0 / 180 * math.pi, car_name="CarIII")
#
#     draw.savefig("car3begin")

def nine_simulation():
    for i in range(2, 11):
        name = str(i - 1)
        name_ = str(i)
        test = CilqrModel(args)
        test.load_NPC_trajectory(f"../trajectory/test{name}.csv")

        test.set_current_state(46 - 4 * (i - 1), 3.4 / 2, 0, 0)
        test.create_global_plan(3.4 / 2, 70 / 3.6)
        test.create_ilqr_agent()
        test.load_restrained_vehicle(f"../trajectory/test{name}.csv")
        desired_path, local_plan, control, X = test.run_step_ilqr()
        print(X)
        test.saveTrajectory(f"../trajectory/test{name_}.csv")

        draw = Draw()
        draw.initRoad()
        draw.drawTrajectory(f"../trajectory/test{name}.csv", "red")
        draw.drawTrajectory(f"../trajectory/test{name_}.csv", "black")
        draw.drawCar(46 - 4 * (i - 1), draw.road_width / 4, 0, 0 / 180 * math.pi, car_name="CarII")

        draw.savefig(f"car{name_}begin")

# nine_simulation()
argparser = argparse.ArgumentParser(description='CARLA CILQR')
add_arguments(argparser)
args = argparser.parse_args()
args.horizon = 40
args.timestep = 0.2
args.horizon = 40
args.x_min = 0.0
args.x_max = 1000000.0
args.y_min = 0.0
args.y_max = 3.4
args.v_min = 0.0
args.v_max = 23.0
args.w_state_bound = 100000000
args.w_yawrate = 10000
args.timestep = 0.2
args.w_vel = 5
args.w_offset = 3
args.offset = 10
args.acc_limits[1] = 4.0

if __name__ == "__main__":

    print("w_val", args.w_vel)

    print("wacc",args.acc_limits)

    test = CilqrModel(args)
    test.load_NPC_trajectory("../trajectory/zero.csv")
    test.set_current_state(46, 3.4 / 2, 0, 0)
    test.create_global_plan(3.4 / 2, 70 / 3.6)
    test.create_ilqr_agent()
    desired_path, local_plan, control, X = test.run_step_ilqr()
    test.saveTrajectory("../trajectory/test1.csv")
    print("desired_path", desired_path)
    print("local_plan", local_plan)
    print("control", control)
    print(type(X))
    print("X", X)
    print("size", len(X))
    nine_simulation()